
from meta_service.store import ArxivPaper

def get_abs(entry_id:str) -> str:
    paper = ArxivPaper(entry_id,False)
    return paper.get_summary()

from meta_service.translate.sync_pool import TranslatePool

# 全局翻译池
GlobalTranslatePool = TranslatePool(
    # proxies={'https': 'http://localhost:7890'},
    frequency_limit=30,
    sentence_limit=2000
)

ABS_JSON_DUMP_PATH = "output/abs_json_dump"

# 摘要翻译
def translate_abs_en2zh(entry_id:str) -> str:
    paper = ArxivPaper(entry_id,False)
    return GlobalTranslatePool.translate_sync(paper.get_summary())

import json
import os

def mkdir_if_not_exist(path:str):
    if not os.path.exists(path):
        os.makedirs(path)

def DumpJson(entry_id:str):
    dump_to_json(entry_id)

def dump_to_json(entry_id:str) -> str:
    # paper = ArxivPaper(entry_id,False)
    paper = ArxivPaper(entry_id)
    title = paper.get_title()
    summary_en = paper.get_summary()
    summary_zh = GlobalTranslatePool.translate_sync(summary_en)
    print("--->",summary_zh)
    entry_id = paper.get_entry_id()
    info = {
        "entry_id":entry_id,
        "title":title,
        "summary_en":summary_en,
        "summary_zh":summary_zh
    }
    mkdir_if_not_exist(ABS_JSON_DUMP_PATH)
    with open(f"{ABS_JSON_DUMP_PATH}/{entry_id}.json", 'w') as f:
        f.write(json.dumps(info, indent=4, ensure_ascii=False))

def show_abs():
    for file_name in os.listdir(ABS_JSON_DUMP_PATH):
        with open(os.path.join(ABS_JSON_DUMP_PATH,file_name), 'r') as f:
            info = json.load(f)
            print(f"entry_id: {info['entry_id']}")
            print(f"title: {info['title']}")
            print(f"summary_en: {info['summary_en']}")
            print(f"summary_zh: {info['summary_zh']}")
            print("\n")
        input("Press Enter to continue...")


if __name__ == "__main__":

    # demo , don't run
    from abs_translate.get_abs import DumpJson,show_abs
    import os
    for file_name in os.listdir("../data/"):
        entry_id = file_name.split(".pkl")[0]
        DumpJson(entry_id)
    show_abs()

